Skip to content

NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator#3798

Open
singhharsh1708 wants to merge 2 commits into
SciML:masterfrom
singhharsh1708:feat/oop-nonlinearsolvealg-wreuse
Open

NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator#3798
singhharsh1708 wants to merge 2 commits into
SciML:masterfrom
singhharsh1708:feat/oop-nonlinearsolvealg-wreuse

Conversation

@singhharsh1708

Copy link
Copy Markdown
Contributor

The OOP NonlinearSolveAlg branch of build_nlsolver builds J, W via build_J_W and then discards both — the NonlinearProblem is built with no jac, and NonlinearSolveCache stores nothing for J/W. Every Newton iter inside the OOP path recomputes the Jacobian via AD/FD with no cross-step reuse.

This is the OOP mirror of #3753 (which fixed the same gap on the IIP path for WOperator with a concrete underlying J). The architecture is identical: detect a reusable W at build time, wire a closure jac into the inner NonlinearProblem, store J/W on the cache so compute_step!'s existing recompute_jacobian flow picks them up.

Scope

Enables W-reuse when W isa Union{WOperator, StaticWOperator} (the case produced for OOP problems with concrete_jac=true, a matrix-free linsolve like KrylovJL_GMRES, or a SciMLOperator jac_prototype). Other W types (plain Matrix, Factorization, generic AbstractSciMLOperator) fall through to the existing no-reuse default — so this is purely additive.

Changes

  1. utils.jl OOP build_nlsolver for NonlinearSolveAlg: detect use_w_reuse, build a closure (z, p) -> convert(AbstractMatrix, W), pass it as jac to NonlinearFunction, store J, W on NonlinearSolveCache.
  2. newton.jl OOP initialize! for NonlinearSolveAlg: when cache.W !== nothing, run the same first-call / new_W_dt_cutoff / divergence logic the IIP path uses, calling update_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma) to refresh the WOperator in place and toggling cache.new_W for compute_step!.

Local validation

Robertson, OOP, ImplicitEuler(linsolve = KrylovJL_GMRES(), concrete_jac = true, ...):

retcode steps nf njacs nw norm(u_end vs NLNewton)
NLNewton Success 1302 2560 641 641
NonlinearSolveAlg Success 1302 2386 0 641 2.3e-13

njacs = 0 on the NS side reflects that the inner solver no longer computes J itself — the WOperator's update_coefficients! does it once per step instead. nw = 641 over 1302 steps means W is reused across ~2 steps on average, matching the NLNewton modified-Newton cadence.

)
prob = if use_w_reuse
nlf_jac = let W = W
(z, p) -> convert(AbstractMatrix, W)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We shouldn't unconditionally convert.

@singhharsh1708

Copy link
Copy Markdown
Contributor Author

good catch — split the closure so StaticWOperator returns its .W field directly, only WOperator path converts now

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch from 350a7e8 to 0f20284 Compare June 29, 2026 18:51
end
else
let Ww = W
(z, p) -> convert(AbstractMatrix, Ww)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about non-static opertators

@singhharsh1708

singhharsh1708 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

for non-static WOperator we shouldn't unconditionally convert either.
Updated to mirror the IIP gate from #3753: use_w_reuse now requires either StaticWOperator or a WOperator whose underlying J is concrete (not an AbstractSciMLOperator). The else branch hands NonlinearSolve the cached Ww._concrete_form directly via the closure (no convert call). Matrix-free / non-concrete WOperator cases fall back to the non-w-reuse path.

end
else
let Ww = W
(z, p) -> Ww._concrete_form

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use the operator...

@singhharsh1708

singhharsh1708 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

fixed — now hands the operator through directly

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch 2 times, most recently from 833ee15 to e414df1 Compare July 3, 2026 19:27
@singhharsh1708

singhharsh1708 commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

Hi @ChrisRackauckas — quick nudge on this stack when you have a moment. All three PRs (#3795, #3796, #3798) have addressed the earlier review comments and are rebased on the latest master:

Once the next SciMLOperators release tags, I'll open a follow-up here to route OOP + SciMLOperator mass matrix through WOperator and drop #3796's informative error for that path.

No rush — just wanted to make sure these weren't waiting on me. Thanks!

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

Ran this branch against an OOP nonlinear problem and hit two defects in the W-reuse path — posting reproductions so they can be addressed:

1. The WOperator jac closure crashes on first Jacobian use (plain Vector problems).

using OrdinaryDiffEqSDIRK, OrdinaryDiffEqNonlinearSolve, NonlinearSolve, ADTypes
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
vdp(u, p, t) = [u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])]
prob = ODEProblem(vdp, [2.0, 0.0], (0.0, 1.0), [1.0e3])
integ = init(prob, TRBDF2(nlsolve = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff())), concrete_jac = true); reltol = 1e-8, abstol = 1e-10)
step!(integ)
# ERROR: MethodError: Cannot `convert` an object of type WOperator{false, ...} to an object of type Matrix{Float64}

The (z, p) -> Ww closure hands the raw WOperator to NonlinearSolve, whose jacobian cache holds a dense Matrix (no jac_prototype on the inner NonlinearFunction), so the assignment tries convert(Matrix, WOperator) and throws. init succeeds — it dies on the first step! that evaluates the jac, which is also why this wasn't caught: the PR currently has no tests. A test exercising exactly the TRBDF2 + concrete_jac=true + Vector u0 path above (plus a StaticArrays one for StaticWOperator) would have flagged it.

2. Even with the convert fixed, the Jacobian values are never refreshed — J is frozen at its build-time state for the entire integration.

The update in initialize! is update_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma). For WOperator that forwards to update_coefficients!(W.J, u, p, t) (SciMLOperators woperator.jl), which is a no-op when W.J is a plain matrix — and the use_w_reuse gate specifically selects for a concrete non-operator W.J. So only gamma ever updates; the J entries stay at their build_nlsolver-time values. (For StaticWOperator it's worse: Ws.W is a fixed matrix, so not even gamma changes.) Newton silently degrades to a chord iteration with J(u₀, t₀) forever — no failure on mildly nonlinear problems, just quietly worse convergence, and on strongly state-dependent J (e.g. the VdP case above) divergence→step-rejection loops.

The IIP counterpart (#3753) avoids this because _update_nlsolvealg_W! explicitly recomputes J via f.jac/jacobian! before jacobian2W!. The OOP path needs the same: an explicit J recompute (the OOP analog is calc_J/jacobian(uf, uprev, integrator)) and a W rebuild — for StaticWOperator, reconstructing StaticWOperator(J - mass_matrix * inv(dtgamma)); for WOperator, refreshing the concrete J and gamma together — inside the should_update branch.

Also worth folding in while touching this: pass jac_prototype matching W's structure to the inner NonlinearFunction — without it NonlinearSolve allocates a dense J, which is the direct cause of crash 1 and also silently densifies sparse problems (I've opened a separate PR fixing that on the IIP path).

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch from e414df1 to 936de56 Compare July 5, 2026 03:42
@singhharsh1708

Copy link
Copy Markdown
Contributor Author

Thanks for the review @ChrisRackauckas-Claude — both defects reproduced and fixed. New commit bd5af110:

Defect 1 (crash on convert): the OOP closure now returns the concrete matrix directly — W_ref[].W for StaticWOperator, convert(AbstractMatrix, W_ref[]) for WOperator — and I also pass jac_prototype matching that shape to the inner NonlinearFunction. NonlinearSolve no longer tries to convert(Matrix, WOperator).

Defect 2 (frozen J): added _update_nlsolvealg_W_oop! (mirrors the IIP _update_nlsolvealg_W!) — inside should_update we now recompute J via f.jac(uprev, p, tstep) or jacobian(uf, uprev, integrator), then either copyto!(W.J, J_new); W.gamma = dtgamma for WOperator or rebuild StaticWOperator(J_new - mass_matrix * inv(dtgamma)) and swap the Ref for StaticWOperator. uf is now stored in the nlcache so jacobian(uf, ...) works on the OOP path.

Ref indirection: the closure captures W_ref instead of W directly, so the StaticWOperator swap in initialize! is visible to NonlinearSolve without rebuilding the outer NonlinearProblem. For WOperator this is just a passthrough since it's already mutable.

Tests: added two regressions at the end of linear_nonlinear_tests.jl covering exactly the bot's two failure modes:

  • Plain Vector VdP + TRBDF2 + concrete_jac = true — used to crash; now succeeds with njacs > 0.
  • SVector VdP variant — same, exercising the StaticWOperator path.

Locally: both new tests pass; existing linear_nonlinear_tests suite still passes.

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch from bd5af11 to d89e8d7 Compare July 6, 2026 04:53
nlf_jac = let W_ref = W_ref
(z, p) -> convert(AbstractMatrix, W_ref[])
end
jac_prototype = copy(convert(AbstractMatrix, W))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no it's still concretizing.

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch 2 times, most recently from c752622 to 9cfc718 Compare July 6, 2026 10:43
Comment on lines +512 to +526
W_ref = use_w_reuse ? Ref{typeof(W)}(W) : nothing
prob = if use_w_reuse
if W isa StaticWOperator
nlf_jac = let W_ref = W_ref
(z, p) -> W_ref[].W
end
NonlinearProblem(
NonlinearFunction{false, SciMLBase.FullSpecialize}(
nlf; jac = nlf_jac, jac_prototype = W.W
),
copy(ztmp), nlp_params
)
else
nlf_jac = let W_ref = W_ref
(z, p) -> convert(AbstractMatrix, W_ref[])

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OOP is purposefully not mutable. We want to make things fold away when static. This will violate that. I think it's simply not possible to do similar W-reuse in the OOP path without violating the principle of making this path non-mutating.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For StaticWOperator the payload is an immutable SMatrix — reassigning nlcache.W[] re-points the Ref cell, the SMatrix itself is never mutated, so fold-away on the value side still holds. The mutation you're objecting to only shows up on the WOperator branch (copyto!(W.J, …)).
And picking NonlinearSolveAlg already forces a Matrix{Float64} jac cache inside construct_jacobian_cache — the OOP path isn't fold-away in that mode regardless of what this PR does. If Static-only reuse is still off the table given that, I'll close.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it's fine if there's a simplediffeq path that fully non-allocates with simplenonlinearsolve. But the conversion stuff is never okay.

@ChrisRackauckas

Copy link
Copy Markdown
Member

On "still concretizing": NonlinearSolveBase.construct_jacobian_cache calls safe_similar(jac_prototype) and allocates a Matrix{Float64} regardless of what the closure returns — so the OOP jac already flows through a mutable buffer once NonlinearSolveAlg is picked. If that's the real blocker, the fix belongs upstream (similar/copyto! for WOperator in SciMLOperators), not here.

That should get fixed upstream. Making everything dense is not a valid answer.

@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch from 9cfc718 to f4cf1d3 Compare July 7, 2026 09:30
@singhharsh1708

Copy link
Copy Markdown
Contributor Author

Dropped WOperator branch. W-reuse only fires on StaticWOperator + AbstractSimpleNonlinearSolveAlgorithm — SMatrix rebuilt, Ref re-pointed, no convert.

@ChrisRackauckas

Copy link
Copy Markdown
Member

Why no WOperator?

Drop the WOperator branch (concretizing/mutation-based).
SimpleNonlinearSolve algs land in NonlinearSolveNoInitCache which has no stats field.
@singhharsh1708 singhharsh1708 force-pushed the feat/oop-nonlinearsolvealg-wreuse branch from b1b82a7 to 141a00b Compare July 8, 2026 20:12
@singhharsh1708

singhharsh1708 commented Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

@ChrisRackauckas Only StaticWOperator gives a real win in OOP: its W::SMatrix is immutable and stack-allocated, so nlcache.W[] = StaticWOperator(J_new - M*inv(γ)) rebuilds for free with no mutation.

For WOperator, refreshing W.J::Matrix needs either copyto!(W.J, J_new) (the mutation you rejected) or nlcache.W[] = WOperator(mm, γ, J_new, u) — the constructor allocates a fresh _concrete_form Matrix + _func_cache Vector each rebuild.

Can add a WOperator branch if you want the Vector-path coverage.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants